package org.nerve.webmagic.utils;

import org.nerve.webmagic.utils.ImageUtil;
import org.openqa.selenium.WebDriver;
import ru.yandex.qatools.ashot.AShot;
import ru.yandex.qatools.ashot.Screenshot;
import ru.yandex.qatools.ashot.shooting.ShootingStrategies;

import javax.imageio.ImageIO;
import java.io.File;
import java.io.IOException;

/**
 * org.nerve.webmagic.spider
 * Created by zengxm on 2016/4/20 0020.
 */
public class ScreenShotTool {
	public static final String FORMAT="PNG";

	private ImageUtil imageUtil;

	public ScreenShotTool(){
		imageUtil=new ImageUtil();
	}

	public void takeScreenshot(final WebDriver driver,String targetPath){
		Screenshot screenshot = new AShot().shootingStrategy(ShootingStrategies.viewportPasting(500))
				.takeScreenshot(driver);

		File file= new File(targetPath);
		try {
			File dir=file.getParentFile();
			if(!dir.exists())
				dir.mkdirs();

			ImageIO.write(
					screenshot.getImage(),
					FORMAT,
					file);
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	/**
	 * 得到截图的base64编码
	 * @return
	 */
	public String taskScreenshot(final WebDriver driver){
		Screenshot screenshot = new AShot().shootingStrategy(ShootingStrategies.viewportPasting(100))
				.takeScreenshot(driver);

		return imageUtil.imgToBase64String(screenshot.getImage(), FORMAT);
	}
}
